#ifndef TCPFILTER_H
#define TCPFILTER_H

#include <pcap.h>
#include <string>
#include <vector>
#include <map>
#include <netinet/in.h>
#include <netinet/ip.h>
#include <netinet/tcp.h>
#include <ctime>
#include <mutex>
#include <unordered_map>
#include <thread>

// TCP包过滤规则结构体
struct FilterRule {
    uint32_t src_ip;        // 源IP地址（网络字节序）
    uint32_t dst_ip;        // 目的IP地址（网络字节序）
    uint16_t src_port;      // 源端口（主机字节序）
    uint16_t dst_port;      // 目的端口（主机字节序）
    uint8_t flags;          // TCP标志位
    bool drop;              // 是否丢弃匹配的包
    bool log;               // 是否记录匹配的包
};

// SYN洪水防御状态
struct SYNFloodState {
    time_t last_syn_time;   // 最后一个SYN包的时间戳
    int syn_count;          // 单位时间内的SYN包计数
    bool blocked;           // 是否已被封锁
    time_t block_time;      // 封锁开始时间
};

class TCPPacketFilter {
private:
    pcap_t* handle;         // libpcap 句柄
    std::vector<FilterRule> rules;
    bool running;
    std::string device;
    
    // SYN洪水防御相关
    mutable std::mutex flood_mutex; // 线程安全锁
    std::unordered_map<uint32_t, SYNFloodState> syn_flood_map; // IP->防御状态
    int syn_threshold;      // SYN包阈值 (单位时间内的最大SYN包数)
    int time_window;        // 时间窗口 (秒)
    int block_duration;     // 封锁持续时间 (秒)
    
    // 检查数据包是否匹配规则
    bool matchRule(const iphdr* ip_hdr, const tcphdr* tcp_hdr) const;
    
    // 获取匹配规则的动作
    FilterRule getMatchedRule(const iphdr* ip_hdr, const tcphdr* tcp_hdr) const;
    
    // 包处理回调函数
    static void packetHandler(u_char *userData, const struct pcap_pkthdr *pkthdr, const u_char *packet);
    
    // SYN洪水防御处理
    bool handleSYNFlood(uint32_t src_ip);
    
public:
    TCPPacketFilter(const std::string& dev);
    ~TCPPacketFilter();
    
    // 添加过滤规则
    void addRule(uint32_t src_ip, uint32_t dst_ip, 
                uint16_t src_port, uint16_t dst_port,
                uint8_t flags, bool drop, bool log = true);
    
    // 初始化libpcap
    bool init();
    
    // 开始捕获数据包
    void startCapture();
    
    // 停止捕获
    void stopCapture();
    
    // 获取可用网络设备
    static void listDevices();
    
    // 设置SYN洪水防御参数
    void setSYNFloodParams(int threshold, int window, int duration) {
        syn_threshold = threshold;
        time_window = window;
        block_duration = duration;
    }
    
    // 获取当前被封锁的IP列表
    std::vector<uint32_t> getBlockedIPs() const;
    
    // 手动封锁IP
    void blockIP(uint32_t ip);
    
    // 手动解封IP
    void unblockIP(uint32_t ip);
};

#endif // TCPFILTER_H
